import rclpy
from rclpy.node import Node
import numpy as np
from nav_msgs.msg import Path
from geometry_msgs.msg import PoseStamped, PointStamped
from tf2_ros import TransformListener
from tf2_ros.buffer import Buffer
import tf2_ros
from rclpy.time import Time
import cv2
from nav_msgs.msg import OccupancyGrid
import heapq
import math

class GlobalPlanner(Node):

    def __init__(self):
        super().__init__('global_planner_node')
        self.declare_parameter('inflation_radius_value', 18)
        self.declare_parameter('goal_radius_value', 0.6)
        # 机器人的半径，目标步长
        self.declare_parameter('robot_r', 0.2)
        # 地图路径
        self.declare_parameter('map_path','map_3v3.pgm')
        
        self.publisher = self.create_publisher(
            OccupancyGrid, 
            '/inflated_map', 
            10)
        self.point_sub = self.create_subscription(
            PoseStamped,
            'goal_pose',
            self.point_callback,
            10
        )
        
        self.inflated_map=None
        self.timer = self.create_timer(1.0, self.timer_callback)
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, self)
        self.path_pub = self.create_publisher(Path, 'path', 50)
        self.map=None

        # 接收地图
        map_path=self.get_parameter('map_path').get_parameter_value().string_value
        map_image = cv2.imread(map_path, cv2.IMREAD_UNCHANGED)
        # 对数组（例如图像）中的每个元素进行按位求反操作
        map_image = np.invert(map_image)  # Invert the image if negate is set to 1
        # 镜像反转，因为图像坐标系的原点在左上，地图在左下
        map_image=cv2.flip(map_image, 0)
        
        occupancy_grid = OccupancyGrid()
        occupancy_grid.header.frame_id = 'map'
        occupancy_grid.info.resolution = 0.05
        occupancy_grid.info.width = map_image.shape[1]
        occupancy_grid.info.height = map_image.shape[0]
        occupancy_grid.info.origin.position.x = -13.2
        occupancy_grid.info.origin.position.y = -44.8
        # occupancy_grid.info.origin.position.y = -10.4
        occupancy_grid.info.origin.position.z = 0.0
        map_data = map_image.flatten()  # Flatten the map image array
        map_data[map_data == 255] = 100
        map_data[map_data == 1] = 0
        # 这边看是-1，实际上值为255
        map_data[map_data == 50] = -1
        occupancy_grid.data = map_data.astype(np.int8).tolist()
        inflation_radius_value=self.get_parameter('inflation_radius_value').get_parameter_value().integer_value
        self.inflated_map = self.inflate_map(occupancy_grid, inflation_radius=inflation_radius_value) # specify the inflation radius here
        # 用于左右摇摆
        self.left_right_flag=1 

    # 1s定时器
    def timer_callback(self):
        if not self.inflated_map:
            return
        # 发布膨胀的地图
        self.publisher.publish(self.inflated_map)
        self.map=np.array(self.inflated_map.data).reshape((self.inflated_map.info.height, self.inflated_map.info.width))

    # A*搜索
    def astar(self, start, end):
        # 给定的起始点
        orgin_start=start
        orgin_end=end

        # 路径规划的openlist和closelist
        open_list_fore = []
        closed_list_fore = set()
        g_fore = dict()

        # 前向搜索
        parent_fore=dict()
        start_x,start_y=start
        end_x,end_y=end

        if start_x < 0 or start_x >= self.map.shape[1] or start_y < 0 or start_y >= self.map.shape[0]:
            self.get_logger().info("start position is out map")
            return None
        
        if end_x < 0 or end_x >= self.map.shape[1] or end_y < 0 or end_y >= self.map.shape[0]:
            self.get_logger().info("end position is out map")
            return None
        # 脱离障碍物
        if self.map[start_y, start_x] != 0:
            parent_fore,start=self.escape_inflation(start)
        
        if self.map[end_y, end_x] != 0:
            parent_back,end=self.escape_inflation(end)
            
        end_x,end_y=end
        if self.map[end_y, end_x] == 100:
            return None
        
        # 向队列里推送A*的起点
        heapq.heappush(open_list_fore, (0, start))
        
        # 如果队列里有数据
        while open_list_fore:
            # 向前搜
            cost_fore, current_node_fore = heapq.heappop(open_list_fore)
            # 如果当前的点在从终点向起点搜的字典里，结束，计算路径
            if current_node_fore==end:
                # 如果要把障碍物里的部分排除，把orgin_end改为end
                return self.extract_path(parent_fore,orgin_start,end)
               
            # 将当前点加入已经搜过的点集
            closed_list_fore.add(current_node_fore)
            
            for neighbor in self.get_neighbors(current_node_fore):
                # 防止重复搜索
                if neighbor in closed_list_fore:
                    continue
                # 计算代价值
                new_cost_fore = cost_fore + self.get_heuristic(neighbor, end)
                # 如果相邻的点不在字典里或代价值小
                if neighbor not in g_fore:
                    g_fore[neighbor] = math.inf
                if new_cost_fore < g_fore[neighbor]:  
                    g_fore[neighbor] = new_cost_fore
                    parent_fore[neighbor] = current_node_fore
                    heapq.heappush(open_list_fore, (new_cost_fore, neighbor))
        return None
    
    # 生成路径
    def extract_path(self, PARENT_fore,start,end):
        """
        extract path from start and goal
        :param s_meet: meet point of bi-direction a*
        :return: path
        """

        # extract path for foreward part
        path_fore = [end]
        s = end
        robot_r_value=self.get_parameter('robot_r').get_parameter_value().double_value
        set_flag=int(robot_r_value/self.inflated_map.info.resolution)
        flag=set_flag
        while True:
            s = PARENT_fore[s]
            flag+=1
            if flag>=set_flag:
                flag=0
                path_fore.append(s)
            if s == start:
                break

        return list(reversed(path_fore))

    # 搜索相邻的四个点
    def get_neighbors(self, node):
        neighbors = []
        x, y = node
        
        directions = [(1, 0), (0, 1), (-1, 0), (0, -1),(1, 1), (-1, 1), (1, -1), (-1, -1)]
        
        for direction in directions:
            new_x, new_y = x + direction[0], y + direction[1]
            if 0 <= new_x < self.map.shape[1] and 0 <= new_y < self.map.shape[0] and self.map[new_y, new_x] == 0:
                neighbors.append((new_x, new_y))
        
        return neighbors

    # 脱离障碍物
    def escape_inflation(self,node):
        escape_parent=dict()
        inflation_radius_value=self.get_parameter('inflation_radius_value').get_parameter_value().integer_value
        # inflation_radius_value=5
        x,y=node
        output_node=node
        min_len=math.inf
        directions = [(1, 0), (0, 1), (-1, 0), (0, -1)]

        for direction in directions:
            parent_buffer=dict()
            for i in range(1,inflation_radius_value+1):
                new_x = x + direction[0]*i
                new_y = y + direction[1]*i
                if new_x>self.map.shape[1]-1 or new_y>self.map.shape[0]-1 or new_y<0 or new_x<0:
                    continue
                parent_buffer[(new_x, new_y)]=(x + direction[0]*(i-1),y + direction[1]*(i-1))
                if self.map[new_y, new_x] != 100:
                    if i<min_len:
                        escape_parent=parent_buffer
                        min_len=i
                        output_node=(new_x, new_y)
                    break
        return escape_parent,output_node

    # 计算当前点的代价值
    def get_heuristic(self, node, end):
        x1, y1 = node
        x2, y2 = end
        return ((x1-x2)*(x1 - x2)) + ((y1-y2)*(y1 - y2))

    # 收到目标点后
    def point_callback(self, msg):
        if self.map is None:
            return
        
        # goal_action=PoseStamped()
        # goal_action.pose.position.x = msg.pose.position.x+msg.pose.orientation.x*self.left_right_flag
        # goal_action.pose.position.y = msg.pose.position.y+msg.pose.orientation.y*self.left_right_flag
        # self.get_logger().info('I heardx: "%d"' % goal_action.pose.position.x)
        # self.get_logger().info('I heardy: "%d"' % goal_action.pose.position.y)

        try:
            transform = self.tf_buffer.lookup_transform(
                "map", "base_link",  # 目标坐标系和源坐标系
                Time())
        # 输出转化矩阵
        except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException):    
            return 0.0,0.0
        
        goal_radius_value=self.get_parameter('goal_radius_value').get_parameter_value().double_value
        if abs(transform.transform.translation.x-msg.pose.position.x)<goal_radius_value and abs(transform.transform.translation.y-msg.pose.position.y)<goal_radius_value:
            # self.get_logger().info('The number is: %d', number)
            return

        path = self.astar((self.world_to_map(transform.transform.translation.x,transform.transform.translation.y)),
                          (self.world_to_map(msg.pose.position.x,msg.pose.position.y)))
        if not path:
            self.get_logger().info("can not get path")
            return
        path_msg = Path()
        path_msg.header.stamp=self.get_clock().now().to_msg()
        path_msg.header.frame_id='map'
        for node in path:
            pose = PoseStamped()
            pose.header.stamp=self.get_clock().now().to_msg()
            pose.header.frame_id='map'
            pose.pose.position.x,pose.pose.position.y= self.map_to_world(node[0],node[1])
            path_msg.poses.append(pose)
        self.path_pub.publish(path_msg)

    # 对地图进行膨胀操作
    def inflate_map(self, map_msg, inflation_radius):
        # 数据类型定义
        inflated_map = OccupancyGrid()
        inflated_map.header = map_msg.header
        inflated_map.info = map_msg.info
        inflated_data = np.array(map_msg.data).reshape((map_msg.info.height, map_msg.info.width))
        inflated_data_copy = inflated_data.copy()
        # 遍历膨胀
        for y in range(map_msg.info.height):
            for x in range(map_msg.info.width):
                if inflated_data[y][x] == 100:
                    # 仅遍历膨胀圆内的点
                    for i in range(max(0, x - inflation_radius), min(map_msg.info.width, x + inflation_radius + 1)):
                        # 计算垂直边界
                        j_start = max(0, y - int((inflation_radius**2 - (i - x)**2)**0.5))
                        j_end = min(map_msg.info.height, y + int((inflation_radius**2 - (i - x)**2)**0.5) + 1)
                        for j in range(j_start, j_end):
                            # 仅标记未标记的点
                            if inflated_data[j][i] != 100:
                                inflated_data_copy[j][i] = 100
        
        inflated_list = []
        for row in inflated_data_copy:
            for val in row:
                inflated_list.append(int(val))

        inflated_map.data = inflated_list
        
        return inflated_map

    # 地图网格坐标系转换到世界坐标系
    def map_to_world(self, x, y):
        mx=x*self.inflated_map.info.resolution+self.inflated_map.info.origin.position.x
        my=y*self.inflated_map.info.resolution+self.inflated_map.info.origin.position.y
        return mx, my
    
    # 世界坐标系转换到地图坐标系
    def world_to_map(self, x, y):
        wx = x-self.inflated_map.info.origin.position.x
        wy = y-self.inflated_map.info.origin.position.y
        mx = int(wx / self.inflated_map.info.resolution)
        my = int(wy / self.inflated_map.info.resolution)
        return mx, my

def main(args=None):
    rclpy.init(args=args)

    global_planner_node = GlobalPlanner()

    rclpy.spin(global_planner_node)

    global_planner_node.destroy_node()
    rclpy.shutdown()

if __name__ == '__main__':
    main()
